/**
 * @file aicl_op_001.cpp
 *
 * Copyright (c) Huawei Technologies Co., Ltd. 2012-2018. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include <iostream>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>

#include <stdio.h>
#include <math.h>
#include <ctime>
#include <chrono>

#include "acl/acl.h"
#include "acl/acl_op_compiler.h"
#include "aicl.h"
#include "utils.h"

aiclDataType inputDataTypeCast = AICL_FLOAT16;
aiclDataType outputDataTypeCast = AICL_FLOAT16;

using namespace std;

int main()
{
    const char *opType = "Conv2D";
    int numInput = 3;
    int numOutput = 1;
    aiclRet ret = 0;
    aiclRet ret_sync = 0;
    const char *workPath = getenv("OPENCANN_TPATH");
    if (workPath == NULL) {
        ERROR_LOG("OPENCANN_TPATH is not set");
        return OPENCANN_FAILED;
    }
    std::string workPath_s = workPath;
    std::string configPath = workPath_s + "/common/test_data/config/init_config.json";
    // aiclInit
    if (aiclInit(configPath.c_str()) != AICL_RET_SUCCESS) {
        ERROR_LOG("aiclInit failed");
    }
    int32_t deviceId_ = 0;
    aiclrtStream stream_;
    if (aiclrtSetDevice(deviceId_) != AICL_RET_SUCCESS) {
        ERROR_LOG("Set device[%d] failed.", deviceId_);
    }
    ret = aiclrtCreateStream(&stream_);


    // input
    std::vector<int64_t> inputShapeCast { 2, 1024, 1024, 3 };
    std::vector<int64_t> inputFilterShapeCast { 6, 3, 3, 3 };
    std::vector<int64_t> outputShapeCast { 2, 1024, 1024, 6 };

    aiclTensorDesc *inputDescCast[numInput];
    aiclTensorDesc *OutputDescCast[numOutput];

    // Create aiclTensorDesc, to describe the shape/format/datatype, etc.
    inputDescCast[0] =
        aiclCreateTensorDesc(inputDataTypeCast, inputShapeCast.size(), inputShapeCast.data(), AICL_FORMAT_NHWC);
    inputDescCast[1] = aiclCreateTensorDesc(inputDataTypeCast, inputFilterShapeCast.size(), inputFilterShapeCast.data(),
        AICL_FORMAT_NCHW);
    inputDescCast[2] = aiclCreateTensorDesc(AICL_DT_UNDEFINED, 0, nullptr, AICL_FORMAT_UNDEFINED);
    OutputDescCast[0] =
        aiclCreateTensorDesc(outputDataTypeCast, outputShapeCast.size(), outputShapeCast.data(), AICL_FORMAT_NHWC);
    // set Conv2D attr
    aiclopAttr *opAttr = aiclopCreateAttr();
    if (opAttr == nullptr) {
        ERROR_LOG("singleOp create attr failed");
        return OPENCANN_FAILED;
    }

    int64_t intList[4] { 1, 1, 1, 1 };
    int attrNum = 4;
    ret = aiclopSetAttrListInt(opAttr, "strides", attrNum, intList);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("singleOp set strides attr failed");
        aiclopDestroyAttr(opAttr);
        return OPENCANN_FAILED;
    }

    ret = aiclopSetAttrListInt(opAttr, "pads", attrNum, intList);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("singleOp set pads attr failed");
        aiclopDestroyAttr(opAttr);
        return OPENCANN_FAILED;
    }

    ret = aiclopSetAttrListInt(opAttr, "dilations", attrNum, intList);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("singleOp set dilations attr failed");
        aiclopDestroyAttr(opAttr);
        return OPENCANN_FAILED;
    }
    INFO_LOG("singleOp set attr success");

    void *xTensorPtr = nullptr;
    void *filterTensorPtr = nullptr;
    void *outTensorPtr = nullptr;
    uint32_t x_tensor_size;
    uint32_t filter_tensor_size;
    uint32_t out_tensor_size = 25165824;
    std::vector<aiclDataBuffer *> in_buffers;
    std::vector<aiclDataBuffer *> out_buffers;

    // input
    std::string input_x_file = workPath_s + "/common/test_data/data/aicl_op_input_x_float16.bin";
    std::string input_filter_file = workPath_s + "/common/test_data/data/aicl_op_input_filter_float16.bin";
    std::string out_file = workPath_s + "/common/test_data/data/aicl_op_out_float16.bin";

    xTensorPtr = Utils::GetDeviceBufferOfFile(input_x_file, x_tensor_size);
    filterTensorPtr = Utils::GetDeviceBufferOfFile(input_filter_file, filter_tensor_size);

    ret = aiclrtMalloc(&outTensorPtr, out_tensor_size, AICL_MEM_MALLOC_HUGE_FIRST);

    aiclDataBuffer *x_tensor_data = aiclCreateDataBuffer(xTensorPtr, x_tensor_size);

    if (x_tensor_data == nullptr) {
        ERROR_LOG("can't create data buffer, create input failed");
        return OPENCANN_FAILED;
    }

    aiclDataBuffer *filter_tensor_data = aiclCreateDataBuffer(filterTensorPtr, filter_tensor_size);
    if (filter_tensor_data == nullptr) {
        ERROR_LOG("can't create data buffer, create input failed");
        return OPENCANN_FAILED;
    }

    aiclDataBuffer *bias_tensor_data = aiclCreateDataBuffer(nullptr, 0);
    aiclDataBuffer *out_tensor_data = aiclCreateDataBuffer(outTensorPtr, out_tensor_size);
    if (out_tensor_data == nullptr) {
        ERROR_LOG("can't create data buffer, create input failed");
        return OPENCANN_FAILED;
    }
    INFO_LOG("create data buffer success");

    in_buffers.push_back(x_tensor_data);
    in_buffers.push_back(filter_tensor_data);
    in_buffers.push_back(bias_tensor_data);
    out_buffers.push_back(out_tensor_data);

    // execute
    ret = aiclopCompileAndExecute(opType, numInput, inputDescCast, in_buffers.data(), numOutput, OutputDescCast,
        out_buffers.data(), opAttr, AICL_ENGINE_SYS, AICL_COMPILE_SYS, NULL, stream_);

    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("execute singleOp API failed, errorCode is %d", static_cast<int32_t>(ret));
        aiclDestroyTensorDesc(inputDescCast[0]);
        aiclDestroyTensorDesc(OutputDescCast[0]);
        return OPENCANN_FAILED;
    }
    INFO_LOG("execute conv2d success");

    ret_sync = aiclrtSynchronizeStream(stream_);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("execute singleOp conv2d sync failed, errorCode is %d", static_cast<int32_t>(ret));
        aiclDestroyTensorDesc(inputDescCast[0]);
        aiclDestroyTensorDesc(OutputDescCast[0]);
        return OPENCANN_FAILED;
    }
    INFO_LOG("execute sync success");

    if (stream_ != nullptr) {
        ret = aiclrtDestroyStream(stream_);
        if (ret != AICL_RET_SUCCESS) {
            ERROR_LOG("destroy stream failed, errorCode = %d", static_cast<int32_t>(ret));
        }
        stream_ = nullptr;
    }
    INFO_LOG("end to destroy stream");

    ret = aiclrtResetDevice(deviceId_);
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("reset device %d failed, errorCode = %d", deviceId_, static_cast<int32_t>(ret));
        return OPENCANN_FAILED;
    }
    INFO_LOG("end to reset device ");

    ret = aiclFinalize();
    if (ret != AICL_RET_SUCCESS) {
        ERROR_LOG("finalize aicl failed, errorCode = %d", static_cast<int32_t>(ret));
        return OPENCANN_FAILED;
    }
    INFO_LOG("end to finalize aicl");
    return OPENCANN_SUCCESS;
}
